package com.bruce.springai.documentReader;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.function.Predicate;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

@Component
public class MyTikaDocumentReader {
    
    /**
     * 处理单个文件
     * @param filePath
     * @return
     */
    public List<Document> loadSingleFile(String filePath) {
        File file = new File(filePath);
        if (!file.isFile()) {
            throw new IllegalArgumentException("Path is not a file: " + filePath);
        }
        return processFile(file);
    }

    /**
     * 加载单个文件，指定只加载哪些文件格式的文件
     * @param filePath
     * @return
     */
    public List<Document> loadSingleFile(String filePath, Predicate<File> filter) {
        File file = new File(filePath);
        if (!file.isFile()) {
            throw new IllegalArgumentException("Path is not a file: " + filePath);
        }
        return filter.test(file) ? processFile(file) : Collections.emptyList();
    }

    /**
     * 只加载这个目录下的所有文件，不包含子目录下的文件
     * @param dirPath
     * @return
     */
    public List<Document> loadDirectory(String dirPath) {
        return loadDirectory(dirPath, false);
    }

    /**
     * 只加载这个目录下的所有文件，不包含子目录下的文件 指定只加载哪些文件格式的文件
     * @param dirPath
     * @param filter
     * @return
     */
    public List<Document> loadDirectory(String dirPath, Predicate<File> filter) {
        return loadDirectory(dirPath, false, filter);
    }

     /**
     * 加载这个目录下的所有文件，包含子目录下的文件
     * @param dirPath
     * @return
     */
    public List<Document> loadDirectoryAll(String dirPath){
        return loadDirectory(dirPath, true);
    }


    /**
     * 加载这个目录下的所有文件，包含子目录下的文件 指定只加载哪些文件格式的文件
     * @param dirPath
     * @param filter
     * @return
     */
    public List<Document> loadDirectoryAll(String dirPath, Predicate<File> filter) {
        return loadDirectory(dirPath, true, filter);
    }


    /**
     * 加载目录下的文件，过滤某个文件格式
     * @param dirPath
     * @param recursive
     * @param filter
     * @return
     */
    private List<Document> loadDirectory(String dirPath, boolean recursive, Predicate<File> filter) {
        File dir = new File(dirPath);
        if (!dir.isDirectory()) {
            throw new IllegalArgumentException("Path is not a directory: " + dirPath);
        }
        return processDirectory(dir, recursive, filter);
    }
    
    private List<Document> loadDirectory(String dirPath, boolean recursive) {
        File dir = new File(dirPath);
        if (!dir.isDirectory()) {
            throw new IllegalArgumentException("Path is not a directory: " + dirPath);
        }
        return processDirectory(dir, recursive);
    }

    // 私有方法：处理单个文件
    private List<Document> processFile(File file) {
        Resource resource = new FileSystemResource(file);
        TikaDocumentReader reader = new TikaDocumentReader(resource);
        return reader.read();
    }

    // 私有方法：处理目录
    private List<Document> processDirectory(File dir, boolean recursive) {
        List<Document> documents = new ArrayList<>();
        File[] files = dir.listFiles();
        if (files != null) {
            for (File file : files) {
                if (file.isFile()) {
                    documents.addAll(processFile(file));
                } else if (recursive && file.isDirectory()) {
                    documents.addAll(processDirectory(file, true));
                }
            }
        }
        return documents;
    }


    private List<Document> processDirectory(File dir, boolean recursive, Predicate<File> filter) {
        List<Document> documents = new ArrayList<>();
        File[] files = dir.listFiles();
        if (files != null) {
            for (File file : files) {
                if (file.isFile() && filter.test(file)) {
                    documents.addAll(processFile(file));
                } else if (recursive && file.isDirectory()) {
                    documents.addAll(processDirectory(file, true, filter));
                }
            }
        }
        return documents;
    }
}

